local tree_node = {}

function tree_node.new(elem)
    local new_node = { name = elem.name, element = elem, children = {} }
    return setmetatable(new_node, {__index = tree_node})
end

function tree_node:print(depth)
    depth = depth or 0
    local spaces = string.rep(' ', depth)
    print(string.format("%s%s", spaces, self.name))
    for _, v in pairs(self.children) do
        v:print(depth+1)
    end
end

function tree_node:add_child(node)
    self.children[node.name] = self.children[node.name] or node
    return self.children[node.name]
end

function tree_node:get_child(name)
    return self.children[name]
end

function tree_node:walk(walker)
    walker(self)
    for _, v in pairs(self.children) do
        v:walk(walker)
    end
end

local tree = {}

function tree.new()
    local root = tree_node.new({name = '<root>'})
    return setmetatable({root = root}, {__index = tree})
end

function tree.new_node(elem)
    return tree_node.new(elem)
end

function tree:add_node(node, parent)
    parent = parent or self.root
    return parent:add_child(node)
end

return tree
